import os
import shutil
from pathlib import Path
from typing import List

def split_folder(source_folder: str, num_splits: int) -> List[str]:
    """
    将源文件夹分成指定数量的子文件夹
    
    Args:
        source_folder: 源文件夹路径
        num_splits: 要分成的份数
        
    Returns:
        List[str]: 包含所有新建子文件夹路径的列表
    """
    # 确保源文件夹存在
    source_path = Path(source_folder)
    if not source_path.exists():
        raise FileNotFoundError(f"源文件夹 {source_folder} 不存在")
        
    # 获取所有文件
    files = [f for f in source_path.iterdir() if f.is_file()]
    
    if not files:
        raise ValueError(f"文件夹 {source_folder} 中没有文件")
        
    # 计算每个子文件夹应包含的文件数
    files_per_split = len(files) // num_splits
    remainder = len(files) % num_splits
    
    # 创建子文件夹
    split_folders = []
    start_idx = 0
    
    for i in range(num_splits):
        # 创建新的子文件夹
        new_folder = source_path.parent / f"{source_path.name}_split_{i+1}"
        new_folder.mkdir(exist_ok=True)
        split_folders.append(str(new_folder))
        
        # 计算当前子文件夹应包含的文件数
        current_split_size = files_per_split + (1 if i < remainder else 0)
        end_idx = start_idx + current_split_size
        
        # 复制文件到新文件夹
        for file in files[start_idx:end_idx]:
            shutil.copy2(file, new_folder)
            
        start_idx = end_idx
        
    print(f"已将 {source_folder} 分成 {num_splits} 份")
    return split_folders

def merge_splits(split_folders: List[str], target_folder: str) -> None:
    """
    将分批的文件夹合并回原始文件夹
    
    Args:
        split_folders: 分批文件夹路径列表
        target_folder: 目标合并文件夹路径
    """
    target_path = Path(target_folder)
    target_path.mkdir(exist_ok=True)
    
    for split_folder in split_folders:
        split_path = Path(split_folder)
        if not split_path.exists():
            raise FileNotFoundError(f"分批文件夹 {split_folder} 不存在")
            
        # 复制所有文件到目标文件夹
        for file in split_path.iterdir():
            if file.is_file():
                shutil.copy2(file, target_path)
                
    print(f"已将分批文件夹合并到 {target_folder}")

if __name__ == "__main__":
    source_folder = "/data/hyq/data/other_pdf"
    num_splits = 3  # 可以根据需要修改分割份数
    try:
        # 分批文件夹
        split_folders = split_folder(source_folder, num_splits)
        print("创建的子文件夹:")
        for folder in split_folders:
            print(f"- {folder}")
            
        # 还原文件夹（可选）
        # restore_folder = source_folder + "_restored"
        # merge_splits(split_folders, restore_folder)
        # print(f"文件已还原到: {restore_folder}")
        
    except Exception as e:
        print(f"错误: {str(e)}")
